/*
@file: transform.cpp
@author: ZZH
@date: 2022-05-05
@info: 变换函数实现文件
*/
#include "libFun.h"
#include "fft.h"

EXPORT void __dft(pFunCallArg_t pArgs, BasicType* output)
{
    int allCalNum = pArgs->allCalNum;
    BasicType* arg0 = static_cast<BasicType*>(pArgs->args[0]);
    pComplex_t pOut = reinterpret_cast<pComplex_t>(output);

    rdft(pOut, arg0, allCalNum);
}

EXPORT void __idft(pFunCallArg_t pArgs, BasicType* output)
{
    int allCalNum = pArgs->allCalNum;
    BasicType* arg0 = static_cast<BasicType*>(pArgs->args[0]);
    pComplex_t pIn = reinterpret_cast<pComplex_t>(arg0);

    irdft(output, pIn, allCalNum);

    for (size_t i = 0;i < allCalNum;i++)
        output[i] /= allCalNum;
}

EXPORT void __fft(pFunCallArg_t pArgs, BasicType* output)
{
    int allCalNum = pArgs->allCalNum;
    BasicType* arg0 = static_cast<BasicType*>(pArgs->args[0]);
    pComplex_t pOut = reinterpret_cast<pComplex_t>(output);

    rfft(pOut, arg0, allCalNum);
}

EXPORT void __ifft(pFunCallArg_t pArgs, BasicType* output)
{
    int allCalNum = pArgs->allCalNum;
    pComplex_t arg0 = static_cast<pComplex_t>(pArgs->args[0]);

    irfft(output, arg0, allCalNum);

    for (size_t i = 0;i < allCalNum;i++)
        output[i] /= allCalNum;
}

EXPORT void cfft(pFunCallArg_t pArgs, BasicType* output)
{
    int allCalNum = pArgs->allCalNum;
    pComplex_t arg0 = static_cast<pComplex_t>(pArgs->args[0]);
    pComplex_t pOut = reinterpret_cast<pComplex_t>(output);
    
    fft(pOut, arg0, allCalNum);
}

EXPORT void icfft(pFunCallArg_t pArgs, BasicType* output)
{
    int allCalNum = pArgs->allCalNum;
    pComplex_t arg0 = static_cast<pComplex_t>(pArgs->args[0]);
    pComplex_t pOut = reinterpret_cast<pComplex_t>(output);

    ifft(pOut, arg0, allCalNum);

    for (size_t i = 0;i < allCalNum;i++)
    {
        pOut[i].real /= allCalNum;
        pOut[i].image /= allCalNum;
    }
}

LibFunction_t funcs[] = {
    LIB_FUNCTION(__dft, 1, .name = "dft"),
    LIB_FUNCTION(__idft, 1, .name = "idft"),
    LIB_FUNCTION(__fft, 1, .name = "fft"),
    LIB_FUNCTION(__ifft, 1, .name = "ifft"),
    LIB_FUNCTION(cfft, 1),
    LIB_FUNCTION(icfft, 1),
    END_OF_LIB
};

register_function_lib(funcs);
